import numpy as np


def load_data(length=32,
              data_dir_path="/mnt/data/oss/g2p/"):
    vocab_file_path = data_dir_path + "/vocab.txt"
    data_file_path = data_dir_path + '/train.txt'
    phone_id_file_path = data_dir_path + "/phone_id_map.txt"

    vocabs = open(vocab_file_path, encoding='utf-8').read().strip().split('\n')

    lines = open(data_file_path, encoding='utf-8').read().strip().split('\n')

    phone_id_text = []
    phone_ids = open(phone_id_file_path, encoding='utf-8').read().strip().split('\n')
    for phone_id in phone_ids:
        split = phone_id.split(' ')
        phone_id_text.append(split[0])

    label = []
    encoder_input = []
    decoder_input = []
    text = []

    # eos = vocabs.index('<eos>')
    # vocabs.append('<bos>')
    for line in lines:
        split = line.split(' ')
        text.append(split[0])

        encoder_input_array = np.zeros(length)
        input_array_read = np.array(list(split[1].split(',')), dtype=int)
        encoder_input_array[0:input_array_read.shape[0]] = input_array_read
        encoder_input_array[input_array_read.shape[0]] = phone_id_text.index('<eos>')
        encoder_input.append(encoder_input_array)

        label_array_read = np.array(list(split[2].split(',')), dtype=int)
        label_array = np.zeros(length)
        label_array[0:label_array_read.shape[0]] = label_array_read
        label_array[label_array_read.shape[0]] = phone_id_text.index('<eos>')
        label.append(label_array)

        decoder_input_array = np.zeros(length)
        decoder_input_array[0] = phone_id_text.index('<bos>')
        decoder_input_array[1:label_array_read.shape[0] + 1] = label_array_read
        decoder_input_array[label_array_read.shape[0] + 1] = phone_id_text.index('<eos>')
        decoder_input.append(decoder_input_array)

    return np.array(encoder_input, dtype=int), np.array(decoder_input, dtype=int), np.array(label, dtype=int), \
           vocabs, phone_id_text, text
